package middleware

import (
	"context"
	"gin-blog/pkg/logging"
	"gin-blog/pkg/tracer"

	"github.com/opentracing/opentracing-go/ext"

	"github.com/uber/jaeger-client-go"

	"github.com/gin-gonic/gin"
	opentracing "github.com/opentracing/opentracing-go"
)

// Tracing creates a gin middleware for distributed tracing
func Tracing() gin.HandlerFunc {
	return func(c *gin.Context) {
		var newCtx context.Context
		var span opentracing.Span

		// Try to extract span context from HTTP headers
		spanCtx, err := opentracing.GlobalTracer().Extract(
			opentracing.HTTPHeaders,
			opentracing.HTTPHeadersCarrier(c.Request.Header))

		if err != nil {
			// Create new root span if no context found
			span, newCtx = opentracing.StartSpanFromContextWithTracer(
				c.Request.Context(),
				tracer.Tracer,
				c.Request.URL.Path)
			logging.Info("Starting new trace: ", c.Request.URL.Path)
		} else {
			// Create child span if context exists
			span, newCtx = opentracing.StartSpanFromContextWithTracer(
				c.Request.Context(),
				tracer.Tracer,
				c.Request.URL.Path,
				opentracing.ChildOf(spanCtx),
				opentracing.Tag{Key: string(ext.Component), Value: "HTTP"},
			)
			logging.Info("Continuing trace: ", c.Request.URL.Path)
		}
		defer span.Finish()

		// Extract and set trace/span IDs
		var traceID, spanID string
		if sc, ok := span.Context().(jaeger.SpanContext); ok {
			traceID = sc.TraceID().String()
			spanID = sc.SpanID().String()
			logging.Info("TraceID: ", traceID, " SpanID: ", spanID)
		}

		// Set tracing info in context
		c.Set("X-Trace-ID", traceID)
		c.Set("X-Span-ID", spanID)
		c.Request = c.Request.WithContext(newCtx)

		// Continue processing
		c.Next()
	}
}
